import os
import random
import torch
import pandas as pd
import numpy as np
from model import TransformerAutoEncoder
from torch.utils.data import TensorDataset, DataLoader
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import MinMaxScaler
from sklearn.metrics import recall_score
from sklearn.linear_model import LogisticRegression
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier
import joblib


# Set env
def seed_torch(seed=42):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed) # if you are using multi-GPU.
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.enabled = True
seed_torch()

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Load data
train_df = pd.read_csv('dataset/cs-training.csv')
median = train_df.median()
train_df = train_df.fillna(median)
test_df = pd.read_csv('dataset/cs-test.csv')
test_df = test_df.fillna(median)

print(train_df.columns)
train = train_df.values[:,1:]
test = test_df.values[:,1:]

X, y = train[:, 1:], train[:, 0]
X_test, y_test = test[:, 1:], test[:, 0]

# Min-max scaler
scaler = MinMaxScaler().fit(X)
X = scaler.transform(X)
X_test = scaler.transform(X_test)

X_train ,X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)

# Load features
train_features = joblib.load('train_features.npy')
train_features_c = joblib.load('train_features_c.npy')
val_features = joblib.load('val_features.npy')
test_features = joblib.load('test_features.npy')

# DecisionTreeClassifier
dt = DecisionTreeClassifier(random_state=42, class_weight={0:1, 1:2.28}, max_depth=8).fit(X_train, y_train)
dt_preds = dt.predict(X_val)
dt_recall_o = recall_score(y_val, dt_preds)
print('dt_recall_o:', dt_recall_o)

dt = DecisionTreeClassifier(random_state=42, class_weight={0:1, 1:5.6}, max_depth=8, max_features='sqrt').fit(train_features, y_train)
dt_preds = dt.predict(val_features)
dt_recall = recall_score(y_val, dt_preds)
print('transformer_dt_recall:', dt_recall)

dt = DecisionTreeClassifier(random_state=42, class_weight={0:1, 1:5.6}, max_depth=8, max_features='sqrt').fit(train_features_c, y_train)
dt_preds = dt.predict(val_features)
dt_recall_c = recall_score(y_val, dt_preds)
print('transformer_dt_recall_c:', dt_recall_c)


with open(f"log/{os.path.basename(__file__)}.log", "a") as f:
    f.write(f"dt_recall_o: {dt_recall_o}, "
            f"transformer_dt_recall: {dt_recall}, "
            f"transformer_dt_recall_c: {dt_recall_c}"
            f"\n")